import numpy as np
import os
from os import path

class DataGenerator:
    @staticmethod
    def get_single_flow_single_link(directory='single_flow_single_link\\network_0'):
        node_data = []
        link_data = []
        flow_data = []
        flow_label = []
        for i in range(512, 1024 + 512 + 1,10):
            node_data.append([[i, 3],
                         [-1, -1],
                         [-1, -1],
                         [-1, -1]])
            link_data.append(
                             # [[0,1,2,-1, -1, -1, -1, -1],
                             #  [1,2,3,-1, -1, -1, -1, -1],
                             #  [0,i,0,-1, -1, -1, -1, -1]]
                             [[0,1,0],
                              [1,2,i],
                              [2,3,0],])
            flow_data.append([[0,1,2,3,-1]])
            flow_label.append([[i]])

        node_data = np.array(node_data, dtype=np.float32)
        link_data = np.array(link_data, dtype=np.float32)
        flow_data = np.array(flow_data, dtype=np.float32)
        flow_label = np.array(flow_label, dtype=np.float32)

        node_data_path = path.join(directory, 'node_data.txt')
        link_data_path = path.join(directory, 'link_data.txt')
        flow_data_path = path.join(directory, 'flow_data.txt')
        flow_label_path = path.join(directory, 'flow_label.txt')
        with open(node_data_path, 'w') as outfile:
            for slice_2d in node_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(link_data_path, 'w') as outfile:
            for slice_2d in link_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(flow_data_path, 'w') as outfile:
            for slice_2d in flow_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(flow_label_path, 'w') as outfile:
            for slice_2d in flow_label:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')

    @staticmethod
    def get_multi_flow_single_link(directory='multi_flow_single_link\\network_0'):
        node_data = []
        link_data = []
        flow_data = []
        flow_label = []
        for i in range(512, 1024 + 512 + 1,10):
            node_data.append([[i*0.4, 4],
                              [i*0.8, 4],
                              [-1, -1],
                              [-1, -1],
                              [-1, -1]])
            link_data.append(
                             # [[0, 1, 2, 3, -1, -1, -1, -1],
                             #  [2, 2, 3, 4, -1, -1, -1, -1],
                             #  [0, 0, i, 0, -1, -1, -1, -1]]
                             [[0,2,0],
                              [1,2,0],
                              [2,3,i],
                              [3,4,0]])
            flow_data.append([[0, 2, 3, 4, -1],
                              [1, 2, 3, 4, -1]])
            flow_label.append([[i*0.4],
                               [i-i*0.4]])

        node_data = np.array(node_data, dtype=np.float32)
        link_data = np.array(link_data, dtype=np.float32)
        flow_data = np.array(flow_data, dtype=np.float32)
        flow_label = np.array(flow_label, dtype=np.float32)

        node_data_path = path.join(directory, 'node_data.txt')
        link_data_path = path.join(directory, 'link_data.txt')
        flow_data_path = path.join(directory, 'flow_data.txt')
        flow_label_path = path.join(directory, 'flow_label.txt')
        with open(node_data_path, 'w') as outfile:
            for slice_2d in node_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(link_data_path, 'w') as outfile:
            for slice_2d in link_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(flow_data_path, 'w') as outfile:
            for slice_2d in flow_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(flow_label_path, 'w') as outfile:
            for slice_2d in flow_label:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')

    @staticmethod
    def get_multi_flow_multi_link(directory='multi_flow_multi_link\\network_0'):
        node_data = []
        link_data = []
        flow_data = []
        flow_label = []
        for i in range(512, 1024 + 512 + 1,10):
            node_data.append([[i*0.3, 7],
                              [i*0.4, 8],
                              [-1, -1],
                              [i*0.4, 7],
                              [i*0.3, 8],
                              [-1, -1],
                              [-1, -1],
                              [-1, -1],
                              [-1, -1]])
            link_data.append(
                            #  [[0, 1, 2, 3, 4, 5, 6, 6],
                            #   [2, 2, 5, 5, 5, 6, 7, 8],
                            #   [0, 0, i*0.6, 0, 0, i, 0, 0]]
                             [[0,2,0],
                              [1,2,0],
                              [2,5,i*0.6],
                              [3,5,0],
                              [4,5,0],
                              [5,6,i],
                              [6,7,0],
                              [6,8,0]])
            flow_data.append([[0, 2, 5, 6, 7],
                              [1, 2, 5, 6, 8],
                              [3, 5, 6, 7, -1],
                              [4, 5, 6, 8, -1]])
            flow_label.append([[(i-((i-i*0.3)/2)-i*0.3)/2],
                               [(i-((i-i*0.3)/2)-i*0.3)/2],
                               [(i-i*0.3)/2],
                               [i*0.3]])

        node_data = np.array(node_data, dtype=np.float32)
        link_data = np.array(link_data, dtype=np.float32)
        flow_data = np.array(flow_data, dtype=np.float32)
        flow_label = np.array(flow_label, dtype=np.float32)

        node_data_path = path.join(directory, 'node_data.txt')
        link_data_path = path.join(directory, 'link_data.txt')
        flow_data_path = path.join(directory, 'flow_data.txt')
        flow_label_path = path.join(directory, 'flow_label.txt')
        with open(node_data_path, 'w') as outfile:
            for slice_2d in node_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(link_data_path, 'w') as outfile:
            for slice_2d in link_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(flow_data_path, 'w') as outfile:
            for slice_2d in flow_data:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')
        with open(flow_label_path, 'w') as outfile:
            for slice_2d in flow_label:
                np.savetxt(outfile, slice_2d, newline='\n', fmt='%f', delimiter=',')
                outfile.write('\n')

if __name__ == '__main__':
    DataGenerator.get_single_flow_single_link()
    DataGenerator.get_multi_flow_single_link()
    DataGenerator.get_multi_flow_multi_link()
